{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "import glob\n",
    "import torch\n",
    "import os\n",
    "from numpy.fft import fft, ifft, fftshift, ifftshift\n",
    "import pickle\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = 9, 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_echo.demodulators.demodulator_classic import DemodulatorClassic\n",
    "from torch_echo.modulators.modulator_neural import ModulatorNeural\n",
    "from torch_echo.modulators.modulator_classic import ModulatorClassic\n",
    "from torch_echo.utils.util_data import bits_to_integers, integers_to_bits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import echoecho.EchoPacketWrapper as EPW\n",
    "import echoecho.DSPUtil as dsp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def db20(x):\n",
    "    return 20 * np.log10(np.abs(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOME = os.environ[\"HOME\"]\n",
    "os.chdir(os.path.join(HOME,'tmp/snr-runs'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get SNR-run map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ber(file):\n",
    "    with open(file, \"r\") as f:\n",
    "        nsamps = int(f.readline().strip().split()[-2])\n",
    "        ber = float(f.readline().strip().split()[-1])\n",
    "    return nsamps, ber"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dirs = glob.glob('*-srn1')  # Only SRN1 recorded IQ\n",
    "gain_dirs = {}\n",
    "gain_bers = {}\n",
    "for d in dirs:\n",
    "    with open(d + \"/tx-gain\", \"r\") as f:\n",
    "        gain = float(f.read().strip())\n",
    "        gain_dirs[gain] = d\n",
    "        _, ber = get_ber(d + \"/results\")\n",
    "        gain_bers[gain] = ber\n",
    "print(gain_dirs)\n",
    "print(gain_bers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gains = sorted(gain_dirs.keys())\n",
    "bers = [gain_bers[g] for g in gains]\n",
    "plt.plot(gains, bers, label='Classic BER')\n",
    "plt.plot([6,16],[.01,.01], linestyle='--', label='1% BER (Training SNR)')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel('TX Gain dB')\n",
    "plt.ylabel('Classic QPSK BER')\n",
    "plt.yscale('log')\n",
    "plt.title(\"Gain vs BER for Classics\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epw = EPW.EchoPacketWrapper(beta_rrc=0.13, corr_repetitions=2, samps_per_symb=2, cfo_freqs=[0.25], cfo_samps=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bps = 2\n",
    "sps = 2\n",
    "pktsz = int(epw.full_packet_length((512 * 2 + 128 + 96) / bps))\n",
    "iqsz = int(pktsz * sps)\n",
    "pktsz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_pkts(iq, step=20000, iqsize=iqsz, pktsize=pktsz):\n",
    "    noise = []\n",
    "    raw_pkts = []\n",
    "    pkts = []\n",
    "    while iq.size > iqsize and len(pkts) < 100:\n",
    "        offset = epw.find_channel_estimate_field(iq[:step], 50)\n",
    "        if offset is not None:\n",
    "            noise.append(iq[:offset])\n",
    "            raw_pkts.append(iq[offset:offset + iqsz])\n",
    "            pkts.append(dsp.rrc_decimate_fft(raw_pkts[-1], 0.13, 2))\n",
    "            iq = iq[offset + iqsz:]\n",
    "        else:\n",
    "            iq = iq[step:]\n",
    "    return raw_pkts, pkts, noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gain_snrs = {}\n",
    "for g in gains:\n",
    "    gdir = gain_dirs[g]\n",
    "    iq = np.fromfile(gdir + \"/iq-in.bin\", dtype=np.complex64)\n",
    "    raws, pkts, noise = extract_pkts(iq)\n",
    "    noise_pwrs = []\n",
    "    body_pwrs = []\n",
    "    for i in range(len(pkts)):\n",
    "        p = pkts[i]\n",
    "        n = noise[i]\n",
    "        noise_pwrs.append(np.median(db20(dsp.rrc_decimate_fft(n, 0.13, 2))))\n",
    "        body_pwrs.append(np.median(db20(p[600:1100])))\n",
    "    snr = np.mean(body_pwrs) - np.mean(noise_pwrs)\n",
    "    gain_snrs[g] = snr\n",
    "print(gain_snrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gain_snrs2 = {}\n",
    "for g in gains:\n",
    "    gdir = gain_dirs[g]\n",
    "    iq = np.fromfile(gdir + \"/iq-in.bin\", dtype=np.complex64)\n",
    "    raws, pkts, noise = extract_pkts(iq)\n",
    "    noise_pwrs = []\n",
    "    body_pwrs = []\n",
    "    for i in range(len(pkts)):\n",
    "        p = raws[i]\n",
    "        n = noise[i]\n",
    "        noise_pwrs.append(np.median(db20(n)))\n",
    "        body_pwrs.append(np.median(db20(p[1200:2200])))\n",
    "    snr = np.mean(body_pwrs) - np.mean(noise_pwrs)\n",
    "    gain_snrs2[g] = snr\n",
    "print(gain_snrs2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(gains, [gain_snrs[g] for g in gains], label='Results')\n",
    "plt.plot(list(range(6,17)), np.arange(6, 17)+0.5, linestyle='--', label='1-to-1')\n",
    "plt.plot([10.5, 10.5], [7, 16], linestyle='--', label='1% BER Gain')\n",
    "plt.annotate(\"({}, {:.1f})\".format(10.5, gain_snrs[10.5]), (10.75, gain_snrs[10.5]-0.25))\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel(\"TX Gain dB\")\n",
    "plt.ylabel(\"RX SNR dB\")\n",
    "plt.title(\"TX Gain vs Achieved SNR @ 1 Samples Per Symbol\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(gains, [gain_snrs2[g] for g in gains], label='Results')\n",
    "plt.plot(list(range(6,17)), list(range(5, 16)), linestyle='--', label='1-to-1')\n",
    "plt.plot([10.5, 10.5], [5, 15], linestyle='--', label='1% BER Gain')\n",
    "plt.annotate(\"({}, {:.1f})\".format(10.5, gain_snrs2[10.5]), (10.75, gain_snrs2[10.5]-0.25))\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel(\"TX Gain dB\")\n",
    "plt.ylabel(\"RX SNR dB\")\n",
    "plt.title(\"TX Gain vs Achieved SNR @ 2 Samples Per Symbol\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot([gain_snrs[g] for g in gains], bers, label='Classic BER')\n",
    "plt.plot([6,16],[.01,.01], linestyle='--', label='1% BER (Training SNR)')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel('SNR dB')\n",
    "plt.ylabel('Classic QPSK BER')\n",
    "plt.yscale('log')\n",
    "plt.title(\"SNR vs BER for Classics\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BER vs Gain vs SNR "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mygains = np.array([6, 10.5, 13, 14.5, 16])\n",
    "gains = np.array(gains)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "snrs = np.array([gain_snrs[g] for g in gains])\n",
    "bers = np.array(bers)\n",
    "msnr, bsnr = np.polyfit(gains, snrs, 1)\n",
    "linear = lambda x: msnr * x + bsnr\n",
    "coeffs = np.polyfit(snrs, np.log(bers), 2)\n",
    "quadr = lambda x: np.exp(coeffs[0] * x ** 2 + coeffs[1] * x + coeffs[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(gains, [gain_snrs[g] for g in gains], label='Results')\n",
    "plt.plot(gains, linear(gains))\n",
    "plt.plot(list(range(6,17)), list(range(5, 16)), linestyle='--', label='1-to-1')\n",
    "plt.plot([10.5, 10.5], [5, 15], linestyle='--', label='1% BER Gain')\n",
    "plt.annotate(\"({}, {:.1f})\".format(10.5, gain_snrs2[10.5]), (10.75, gain_snrs2[10.5]-0.25))\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel(\"TX Gain dB\")\n",
    "plt.ylabel(\"RX SNR dB\")\n",
    "plt.title(\"TX Gain vs Achieved SNR @ 2 Samples Per Symbol\")\n",
    "plt.show()\n",
    "\n",
    "plt.plot([gain_snrs[g] for g in gains], bers, label='Classic BER')\n",
    "plt.plot(np.arange(6, 16, 0.5), quadr(np.arange(6, 16, 0.5)))\n",
    "plt.plot([6,15],[.01,.01], linestyle='--', label='1% BER (Training SNR)')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel('SNR dB')\n",
    "plt.ylabel('Classic QPSK BER')\n",
    "plt.yscale('log')\n",
    "plt.title(\"SNR vs BER for Classics\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mysnrs = linear(mygains)\n",
    "print(mysnrs)\n",
    "mybers = quadr(mysnrs)\n",
    "print(mybers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = {'gain-snr-linear': [msnr, bsnr],\n",
    "       'snr-ber-logquadratic': list(coeffs)}\n",
    "with open(\"gain-snr-ber-coeffs.pkl\", \"wb\") as f:\n",
    "    pickle.dump(out, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
